#ifndef _HNS_ROCE_PEER_MEM_H
#define _HNS_ROCE_PEER_MEM_H

#include <linux/module.h>
#include <linux/init.h>
#include <linux/slab.h>
#include <linux/errno.h>
#include <linux/export.h>
#include <linux/scatterlist.h>
#include <rdma/ib_umem.h>

#define IB_PEER_MEMORY_NAME_MAX     64
#define IB_PEER_MEMORY_VER_MAX      16

struct hns_umem {
	struct ib_umem umem;

	int page_shift;
	/* peer memory that manages this umem */
	struct ib_peer_memory_client *ib_peer_mem;
	/* peer memory private context */
	void *peer_mem_client_context;
};

struct peer_memory_client {
	char	name[IB_PEER_MEMORY_NAME_MAX];
	char	version[IB_PEER_MEMORY_VER_MAX];
	int (*acquire)(unsigned long addr, size_t size, void *peer_mem_private_data,
		       char *peer_mem_name, void **client_context);
	int (*get_pages)(unsigned long addr,
			 size_t size, int write, int force,
			 struct sg_table *sg_head,
			 void *client_context, u64 core_context);
	void (*put_pages)(struct sg_table *sg_head, void *client_context);
	int (*dma_map)(struct sg_table *sg_head, void *client_context,
		       struct device *dma_device, int *nmap);
	int (*dma_unmap)(struct sg_table *sg_head, void *client_context,
			 struct device  *dma_device);
	unsigned long (*get_page_size)(void *client_context);
	void (*release)(void *client_context);
	void *(*get_context_private_data)(u64 peer_id);
	void (*put_context_private_data)(void *context);
};

struct ib_peer_memory_statistics {
	atomic64_t num_alloc_mrs;
	atomic64_t num_dealloc_mrs;
	atomic64_t num_reg_pages;
	atomic64_t num_dereg_pages;
	atomic64_t num_reg_bytes;
	atomic64_t num_dereg_bytes;
	unsigned long num_free_callbacks;
};

struct ib_ucontext;

struct ib_peer_memory_client {
	const struct peer_memory_client *peer_mem;
	struct list_head	core_peer_list;
	int invalidation_required;
	struct kref ref;
	struct completion unload_comp;
	/* lock is used via the invalidation flow */
	struct mutex lock;
	struct list_head   core_ticket_list;
	u64	last_ticket;
	struct kobject *kobj;
	struct attribute_group peer_mem_attr_group;
	struct ib_peer_memory_statistics stats;
};

static inline struct hns_umem *to_hns_umem(struct ib_umem *umem)
{
	return container_of(umem, struct hns_umem, umem);
}

struct ib_umem* hns_umem_get(struct ib_device *device, unsigned long addr,
			size_t size, int access, void* peer_mem_private_data, char* peer_mem_name);

typedef int (*invalidate_peer_memory)(void *reg_handle, u64 core_context);

void *ib_register_peer_memory_client(const struct peer_memory_client *peer_client,
			invalidate_peer_memory *invalidate_callback);
void ib_unregister_peer_memory_client(void *reg_handle);

struct ib_peer_memory_client *ib_get_peer_client(unsigned long addr, size_t size, void *peer_mem_private_data,
			char *peer_mem_name, void **peer_client_context);
void ib_put_peer_client(struct ib_peer_memory_client *ib_peer_client,
			void *peer_client_context);

void hns_peer_umem_release(struct ib_umem* umem);

#endif